from myutils.helper import get_item_append,get_accuracy_one,tqdm_update_epoch_batch_info


import gc
import torch
from tqdm import tqdm
import numpy as np
from itertools import chain





def train_one(dataLoader,model,loss,optimizer,scheduler,ema,is_test:bool,print_margin_batch:int,epoch_info:str,device,
        list_loss,list_accuracy_Y,list_accuracy_location):
    """单网络训练一个epoch

    Args:
        dataLoader:数据加载器
        model:模型
        loss:损失函数
        optimizer:优化器
        is_test:是否在测试（孪生姊妹网络中不能测试）
        print_margin_batch:每多少batch输出1次，暂不使用
        epoch_info:epoch信息
        
        list_history_one:单网络相关历史,[list_loss_one,list_loss_Y,list_loss_location,
                list_accuracy_Y,list_accuracy_location]

        # list_history_siamese:孪生姊妹相关历史,[list_loss_siamese,list_loss_contrastive]
    """
    


    gc.collect()
    torch.cuda.empty_cache()
    
    list_Y_hat = []
    list_Y = []
    
    tqdm_dataLoader = tqdm(dataLoader,desc=epoch_info)
    # 批次,数据集
    for batch_i, (img,shape,Y,location) in enumerate(tqdm_dataLoader,1):

        # cuda: 垃圾收集，清理缓存，device转cuda
        gc.collect()
        torch.cuda.empty_cache()
        (img  ,shape  ,Y  ,location)= (img.to(device), shape.to(device), Y.to(device), location.to(device))
        
        # 正向传播
        _,location_hat,Y_hat = model(img,shape)
        # 计算loss
        l = loss(Y_hat  ,Y ,location_hat,location)
        
        if is_test:
            list_Y_hat+=[Y_hat.cpu().detach().numpy()]
            list_Y    +=[Y    .cpu().detach().numpy()]
        if not is_test:
            optimizer.zero_grad() # 梯度置为0
            l.backward()   # 反向传播
            optimizer.step()      # 梯度下降
            if scheduler: scheduler.step()
            if ema: ema.update()

        # if is_test:
        #     print()
        #     print(Y_hat.argmax(axis=1))
        #     print(Y.reshape(-1))


        accuracy_Y        = get_accuracy_one(Y_hat,Y)
        accuracy_location = get_accuracy_one(location_hat,location)

        # 转原始类型(one):loss,accuracy
        get_item_append(list_loss,l)
        get_item_append(list_accuracy_Y,accuracy_Y)
        get_item_append(list_accuracy_location,accuracy_location)

        if batch_i%print_margin_batch==0:
            batch_info=f'{batch_i}/{len(dataLoader)}'
            tqdm_update_epoch_batch_info(tqdm_dataLoader,epoch_info,batch_info,print_margin_batch,
                    list_loss,list_accuracy_Y,list_accuracy_location)
    
    return np.array(list(chain(*list_Y))),np.array(list(chain(*list_Y_hat)))